from __future__ import division, print_function, absolute_import

import argparse
import os

import cv2
import numpy as np

import matplotlib.pyplot as plt


# dataset_id = 2
# a_dets = np.loadtxt('./tmp/MOT-02-det.txt')
# # b_dets = np.loadtxt('./deep-sort/MOT16/train/MOT16-02/det/det.txt',delimiter=',')
# b_dets = np.loadtxt('./deep-sort/MOT16/train/MOT16-02/gt/gt.txt',delimiter=',')
# c_dets = np.load('/Users/gerrie/prj/end2end_mot/deep-sort/resources/detections/MOT16_POI_train/MOT16-02.npy')
# # b_dets = np.loadtxt('./deep-sort/MOT16/train/MOT16-02/gt/gt.txt',delimiter=',')
# c_dets = c_dets[:, :10]
def visualize_det(frame_id):
    # frame_id = 10
    boxes_a = a_dets[a_dets[:, 0] == frame_id]
    boxes_b = b_dets[b_dets[:, 0] == frame_id]
    boxes_c = c_dets[c_dets[:, 0] == frame_id]

    img_dir = 'deep-sort/MOT16/train_v1/MOT16-%02d/img1/' % dataset_id
    img = cv2.imread(os.path.join(img_dir, '%06d.jpg' % frame_id))

    for det in boxes_a:
        fid,tid, x, y, w, h = det[:6].astype(np.int32)
        cv2.rectangle(img, (x, y), (x+w, y+h), (255, 0, 0), 1)  # blue my pose result.

    for det in boxes_b:
        fid,tid, x, y, w, h = det[:6].astype(np.int32)
        cv2.rectangle(img, (x, y), (x+w, y+h), (0, 255, 0), 1)  # green

    for det in boxes_c:
        fid,tid, x, y, w, h = det[:6].astype(np.int32)
        cv2.rectangle(img, (x, y), (x+w, y+h), (0, 0, 255), 1)  # poi

    cv2.imshow('test', img)

def json_visualize():
    import json
    # cmu format
    json_path = '/Users/gerrie/prj/end2end_mot/alphapose_tf_results/results_torch/alphapose-results.json'
    img_path = '/Users/gerrie/prj/end2end_mot/alphapose_tf_results/results_torch/vis/1.jpg'
    frames = json.load(open(json_path, 'r'))
    identity = frames['1.jpg']['bodies'][0]
    joints = identity['joints']
    img = cv2.imread(img_path)
    for i in range(len(joints)//3):
        x, y = joints[3*i], joints[3*i+1]
        cv2.putText(img, '%d' %i, tuple(map(int, (x, y))), fontFace=cv2.FONT_HERSHEY_COMPLEX, fontScale=.5, color=(0, 0, 255))
    cv2.imshow('t', img)
    cv2.waitKey(0)


def main():
    for i in range(1, 200):
        visualize_det(i)
        if cv2.waitKey() == ord('q'):
            break

if __name__ == '__main__':
    # main()
    json_visualize()